import time
from dataclasses import dataclass

import numpy as np

from robots.piper import PiperArm


@dataclass
class ArmStatus:

    is_moving: bool = False
    gripping_sth: bool = False
    effector_xeuler: np.ndarray = None
    target_xeuler: np.ndarray = None


class SkilledPiper(PiperArm):
    """
    A piper arm with some pre-defined skills
    """

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        self.skill_set = {
            "raise_to",
            "pick_lift",
            "place_down",
            "press_hold",
            "pick_place",
            "line_move",
            "swing_forward",
        }

    def raise_to(self, height: float):
        xeuler = self.get_effector_xeuler()
        self.move_to(np.append(xeuler[:2], height), xeuler[3:])

    def pick_lift(self, xyz: np.ndarray, height: float = 0.15, num_attempts: int = 1):
        """
        Lift the object up and retry to grasp the object if the gripper detects no object.
        """
        self.open_gripper()
        self.move_to(xyz)
        self.close_gripper()
        self.raise_to(height)

        attempt = 0
        while not self.is_gripping_sth() and attempt < num_attempts:
            self.open_gripper()
            # compute the new pick position
            # this could be improved by using a more accurate method, say, an adjustment model or camera re-planning
            pick_pos = np.concatenate([xyz[:2], [xyz[2] - 0.005 * (attempt + 1)]])
            self.move_to(pick_pos)
            self.close_gripper()
            self.raise_to(height)
            time.sleep(0.2)
            attempt += 1

        return self.is_gripping_sth()

    def place_down(self, xyz: np.ndarray, height: float = None, release: bool = True):
        """
        If can be used to press and hold the object by setting open_gripper=False
        """
        self.close_gripper()
        if height is not None:
            xyz = np.append(xyz[:2], height)
        self.move_to(xyz)
        if release:
            self.open_gripper()

    def press_hold(self, xyz: np.ndarray):
        self.place_down(xyz, height=0, release=False)
        return self.is_in_motion()  # whether the arm is holding the object

    def pick_place(self, xyz: np.ndarray, end_xyz: np.ndarray, height: float = 0.15):
        if self.pick_lift(xyz, height=height):
            start_pos = self.get_effector_xeuler()[:3]
            end_pos = np.append(end_xyz[:2], start_pos[2])
            self.line_move(start_pos, end_pos)
            self.place_down(end_xyz)
        else:
            print("Failed to grasp the object!")
            return False
        return True

    def line_move(self, xyz: np.ndarray, end_xyz: np.ndarray, height: float = None, num_steps: int = 50):
        assert num_steps >= 2, "num_steps must be at least 2"
        if height is not None:
            xyz = np.append(xyz[:2], height)
            end_xyz = np.append(end_xyz[:2], height)
        self.follow_traj(np.linspace(xyz, end_xyz, num_steps))

    def swing_forward(self, amplitude: float = 0.1, angle: float = np.pi / 1.5):
        xeuler = self.get_effector_xeuler()
        dz = amplitude * np.sin(angle)
        dx = amplitude * np.cos(angle)
        high_xyz = xeuler[:3] + np.array([dx, 0, dz])
        # moving forward & backward
        with self.no_wait_execution():
            self.line_move(xeuler[:3], high_xyz, num_steps=2)
            time.sleep(0.2)
            self.line_move(high_xyz, xeuler[:3], num_steps=2)

    def get_status(self):
        """
        Can be more detailed
        """
        return ArmStatus(
            is_moving=self.is_in_motion(),
            gripping_sth=self.is_gripping_sth(),
            effector_xeuler=self.get_effector_xeuler(),
            target_xeuler=self.ik.get_target_xeuler(),
        )


if __name__ == "__main__":
    self = SkilledPiper("can0", wait_execution=True)
    pos1 = np.array([0.408, 0.168, 0.045])
    pos2 = np.array([0.508, -0.168, 0.045])

    self.raise_to(0.6)
